// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/substrait/substrait_parser.h>
#include <string>
#include <pollux/common/base/exceptions.h>
#include <pollux/substrait/type_utils.h>
#include <pollux/substrait/pollux_substrait_signature.h>

namespace kumo::pollux::substrait {
    TypePtr SubstraitParser::parseType(const ::substrait::Type &substraitType) {
        switch (substraitType.kind_case()) {
            case ::substrait::Type::KindCase::kBool:
                return BOOLEAN();
            case ::substrait::Type::KindCase::kI8:
                return TINYINT();
            case ::substrait::Type::KindCase::kI16:
                return SMALLINT();
            case ::substrait::Type::KindCase::kI32:
                return INTEGER();
            case ::substrait::Type::KindCase::kI64:
                return BIGINT();
            case ::substrait::Type::KindCase::kFp32:
                return REAL();
            case ::substrait::Type::KindCase::kFp64:
                return DOUBLE();
            case ::substrait::Type::KindCase::kString:
                return VARCHAR();
            case ::substrait::Type::KindCase::kBinary:
                return VARBINARY();
            case ::substrait::Type::KindCase::kStruct: {
                const auto &substraitStruct = substraitType.struct_();
                const auto &structTypes = substraitStruct.types();
                std::vector<TypePtr> types;
                std::vector<std::string> names;
                for (int i = 0; i < structTypes.size(); i++) {
                    types.emplace_back(parseType(structTypes[i]));
                    names.emplace_back(fmt::format("col_{}", i));
                }
                return ROW(std::move(names), std::move(types));
            }
            case ::substrait::Type::KindCase::kList: {
                // The type name of list is in the format of: ARRAY<T>.
                const auto &fieldType = substraitType.list().type();
                return ARRAY(parseType(fieldType));
            }
            case ::substrait::Type::KindCase::kMap: {
                // The type name of map is in the format of: MAP<K,V>.
                const auto &sMap = substraitType.map();
                const auto &keyType = sMap.key();
                const auto &valueType = sMap.value();
                return MAP(parseType(keyType), parseType(valueType));
            }
            case ::substrait::Type::KindCase::kUserDefined:
                // We only support UNKNOWN type to handle the null literal whose type is
                // not known.
                return UNKNOWN();
            case ::substrait::Type::KindCase::kDate:
                return DATE();
            default:
                POLLUX_NYI(
                    "Parsing for Substrait type not supported: {}",
                    substraitType.DebugString());
        }
    }

    std::vector<TypePtr> SubstraitParser::parseNamedStruct(
        const ::substrait::NamedStruct &namedStruct) {
        // Nte that "names" are not used.

        // Parse Struct.
        const auto &substraitStruct = namedStruct.struct_();
        const auto &substraitTypes = substraitStruct.types();
        std::vector<TypePtr> typeList;
        typeList.reserve(substraitTypes.size());
        for (const auto &type: substraitTypes) {
            typeList.emplace_back(parseType(type));
        }
        return typeList;
    }

    int32_t SubstraitParser::parseReferenceSegment(
        const ::substrait::Expression::ReferenceSegment &refSegment) {
        auto typeCase = refSegment.reference_type_case();
        switch (typeCase) {
            case ::substrait::Expression::ReferenceSegment::ReferenceTypeCase::
            kStructField: {
                return refSegment.struct_field().field();
            }
            default:
                POLLUX_NYI(
                    "Substrait conversion not supported for ReferenceSegment '{}'",
                    std::to_string(typeCase));
        }
    }

    std::vector<std::string> SubstraitParser::makeNames(
        const std::string &prefix,
        int size) {
        std::vector<std::string> names;
        names.reserve(size);
        for (int i = 0; i < size; i++) {
            names.emplace_back(fmt::format("{}_{}", prefix, i));
        }
        return names;
    }

    std::string SubstraitParser::makeNodeName(int node_id, int col_idx) {
        return fmt::format("n{}_{}", node_id, col_idx);
    }

    int SubstraitParser::getIdxFromNodeName(const std::string &nodeName) {
        // Get the position of "_" in the function name.
        std::size_t pos = nodeName.find("_");
        if (pos == std::string::npos) {
            POLLUX_FAIL("Invalid node name.");
        }
        if (pos == nodeName.size() - 1) {
            POLLUX_FAIL("Invalid node name.");
        }
        // Get the column index.
        std::string colIdx = nodeName.substr(pos + 1);
        try {
            return stoi(colIdx);
        } catch (const std::exception &err) {
            POLLUX_FAIL(err.what());
        }
    }

    const std::string &SubstraitParser::findFunctionSpec(
        const melon::F14FastMap<uint64_t, std::string> &functionMap,
        uint64_t id) const {
        if (functionMap.find(id) == functionMap.end()) {
            POLLUX_FAIL("Could not find function id {} in function map.", id);
        }
        melon::F14FastMap<uint64_t, std::string> &map =
                const_cast<melon::F14FastMap<uint64_t, std::string> &>(functionMap);
        return map[id];
    }

    std::string SubstraitParser::findPolluxFunction(
        const melon::F14FastMap<uint64_t, std::string> &functionMap,
        uint64_t id) const {
        std::string funcSpec = findFunctionSpec(functionMap, id);
        std::string_view funcName = getNameBeforeDelimiter(funcSpec, ":");
        return mapToPolluxFunction({funcName.begin(), funcName.end()});
    }

    std::string SubstraitParser::mapToPolluxFunction(
        const std::string &substraitFunction) const {
        auto it = substraitPolluxFunctionMap_.find(substraitFunction);
        if (it != substraitPolluxFunctionMap_.end()) {
            return it->second;
        }

        // If not finding the mapping from Substrait function name to Pollux function
        // name, the original Substrait function name will be used.
        return substraitFunction;
    }

    std::vector<std::string> SubstraitParser::getSubFunctionTypes(
        const std::string &substraitFunction) {
        // Get the position of ":" in the function name.
        size_t pos = substraitFunction.find(":");
        // Get the input types.
        std::vector<std::string> types;
        if (pos == std::string::npos || pos == substraitFunction.size() - 1) {
            return types;
        }
        // Extract input types with delimiter.
        for (;;) {
            const size_t endPos = substraitFunction.find("_", pos + 1);
            if (endPos == std::string::npos) {
                std::string typeName = substraitFunction.substr(pos + 1);
                if (typeName != "opt" && typeName != "req") {
                    types.emplace_back(typeName);
                }
                break;
            }

            const std::string typeName =
                    substraitFunction.substr(pos + 1, endPos - pos - 1);
            if (typeName != "opt" && typeName != "req") {
                types.emplace_back(typeName);
            }
            pos = endPos;
        }
        return types;
    }

    std::vector<TypePtr> SubstraitParser::getInputTypes(
        const std::string &signature) {
        std::vector<std::string> typeStrs = getSubFunctionTypes(signature);
        std::vector<TypePtr> types;
        types.reserve(typeStrs.size());
        for (const auto &typeStr: typeStrs) {
            types.emplace_back(
                PolluxSubstraitSignature::fromSubstraitSignature(typeStr));
        }
        return types;
    }
} // namespace kumo::pollux::substrait
